#!/usr/bin/python
# -*- coding:utf-8 -*-

'''This tool read all sessions from register table of putty
then list them to stdout. 
'''
import sys
import urllib
import _winreg as reg
import traceback
import argparse #this need 2.7+

PUTTY_SESSION_KEY_STR=u'Software\SimonTatham\PuTTY\Sessions'
DEFAULT_SETTINGS=urllib.quote(u'Default Settings')

def getargparser():
    parser = argparse.ArgumentParser(description='List putty sessions.')
    parser.add_argument('-l', '--list',
                        action='store_true',
                        help='list default output, same as -o Name Protocol AgentFwd LocalUserName')
    parser.add_argument('-o', '--output',
                        action='append',
                        nargs='+',
                        help='the names of option will be outputted, see [--listoptions] for all available options',
                        metavar='OptionName')
    parser.add_argument('-v', '--invert-match',
                        action='append',
                        nargs='+',
                        help='the names of session will be omitted. Like grep -v.',
                        metavar='OmittedSessionName',
                        dest='omits')
    parser.add_argument('--listoptions',
                        action='store_true',
                        help='list all available options.')
    parser.add_argument('sessions',
                        action='append',
                        nargs='*',
                        help='session names to list.',
                        metavar='Session')
    
    args = parser.parse_args()
    if args.output:
        map(args.output[0].extend, args.output[1:])
        args.output = args.output[0]
    else:
        args.output = []
    if args.omits:
        map(args.omits[0].extend, args.omits[1:])
        args.omits = args.omits[0]
    else:
        args.omits = []
    if args.sessions:
        map(args.sessions[0].extend, args.sessions[1:])
        args.sessions = args.sessions[0]
    else:
        args.sessions = []
    if args.list and args.output:
        parser.error('[-l] and [-o] are exclusive\n')
    if args.list:
        args.output = ['Protocol', 'AgentFwd', 'LocalUserName']
    
    return args

class Session(object):
    def str_me(self, output):
        l = [self.Name]
        l.extend(str(getattr(self, name)) for name in output)
        return '\t'.join(l)
    
def main():
    args=getargparser()
    k=reg.CreateKey(reg.HKEY_CURRENT_USER, PUTTY_SESSION_KEY_STR)
    
    # the --listoptions
    if args.listoptions:
        skey = reg.CreateKey(k, DEFAULT_SETTINGS)
        options=[]
        try:
            index=0
            while True:
                options.append(reg.EnumValue(skey, index)[0])
                index+=1
        except WindowsError,e:
            if e.winerror != 259:
                raise
        options.sort()
        print '\n'.join(options)
        return 0
    # other args
    if args.sessions:
        sessionnames=args.sessions
        #check existence
        for sessionname in sessionnames:
            try:
                reg.OpenKey(k, sessionname)
            except WindowsError,e:
                if e.winerror == 2:
                    raise Exception('Session %s does not exists.' % sessionname)
                else:
                    raise
    else:
        sessionnames = []
        try:
            index=0
            while True:
                sessionnames.append(reg.EnumKey(k,index))
                index += 1
        except WindowsError,e:
            if e.winerror != 259:
                raise
        
    s=[]        #list of Session
    for skeyname in sessionnames:
        session = Session()
        session.Name = urllib.unquote(skeyname)
        skey = reg.CreateKey(k, skeyname)
        for name in args.output:
            try:
                value = reg.QueryValueEx(skey, name)[0]
            except WindowsError,e:
                if e.winerror == 2:
                    raise Exception('Option %s does not exists. see --listoptions for available options.' % name)
                else:
                    raise
            setattr(session, name, value)
        s.append(session)
            
    s=filt_session(s, args.omits)
    for session in s:
        sys.stdout.write(session.str_me(args.output)+'\n')
    return 0

def filt_session(sessions, omits):
    if not omits:
        return sessions
    return list(session for session in sessions if not session.Name in omits)

if __name__ == "__main__":
    sys.exit(main())
    

